{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d006b2ea-9dfe-49c7-88a9-a5a0775185fd",
   "metadata": {},
   "source": [
    "# A tool to evaluate a mathematical expression\n",
    "\n",
    "This week the tool used in FlightAI was a database lookup function.\n",
    "\n",
    "Here I implement a python code interpreter function as tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b0e8691-71f9-486c-859d-ea371401dfa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e2792ae-ff53-4b83-b2c3-866533ba2b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load environment variables in a file called .env\n",
    "# Print the key prefixes to help with any debugging\n",
    "\n",
    "load_dotenv()\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')\n",
    "google_api_key = os.getenv('GOOGLE_API_KEY')\n",
    "\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "if anthropic_api_key:\n",
    "    print(f\"Anthropic API Key exists and begins {anthropic_api_key[:7]}\")\n",
    "else:\n",
    "    print(\"Anthropic API Key not set\")\n",
    "\n",
    "if google_api_key:\n",
    "    print(f\"Google API Key exists and begins {google_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"Google API Key not set\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79e44ee9-af02-448c-a747-17780ee55791",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai = OpenAI()\n",
    "MODEL = \"gpt-4o-mini\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33ec55b1-0eff-43f1-9346-28145fa2fc47",
   "metadata": {},
   "source": [
    "# Defining the tool function\n",
    "\n",
    "Add print statements to make sure the function is used instead of the native gpt interpreter capability.\n",
    "\n",
    "I used multi shot in the system prompt to make sure gpt generate the code in the format that the tool accept."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94e0e171-4975-457b-88cb-c0d90f51ca65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_math_expression(my_code):\n",
    "    print(f\"EXECUTING FUNCTION WITH CODE: {my_code}\")\n",
    "    exec(my_code)\n",
    "    r = locals()['interpreter_result'] \n",
    "    return r\n",
    "\n",
    "\n",
    "math_function = {\n",
    "    \"name\": \"evaluate_math_expression\",\n",
    "    \"description\": \"Give the result of a math expression. \\\n",
    "    Call this whenever you need to know the result of a mathematical expression. \\\n",
    "    Generate python code ALWAYS with the final result assigned to a variable called 'interpreter_result'. \\\n",
    "    For example when a user asks 'What is 2+2' generate 'interpreter_result = 2+2', and pass this code to the tool. \\\n",
    "    Another example if a user ask 'What is log(5)' generate 'import math; interpreter_result = math.log(5)' and pass this code to the tool.\",\n",
    "    \n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"my_code\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The python math expression to evaluate\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"my_code\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}\n",
    "\n",
    "tools = [{\"type\": \"function\", \"function\": math_function}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c85c01cc-776e-4a9d-b506-ea0d68fc072d",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_math_expression(\"import math; interpreter_result = math.log(5)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "858c5848-5835-4dff-9dc0-68babd367e11",
   "metadata": {},
   "source": [
    "# Using the tool in a UI program\n",
    "\n",
    "You can ask messages like:\n",
    "- \"What is 2+2?\"\n",
    "- \"What is 3 power 2?\"\n",
    "- \"I have 25 apples. I buy 10 apples. How manny apples do I have?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c119b48b-d4b4-41ae-aa2f-2ec2f09af2f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"You are a math assistant. \\\n",
    "Generate python code to give result of a math expression, always name the result 'interpreter_result'. \\\n",
    "For example when a user asks 'What is 2+2', generate 'interpreter_result = 2+2' and pass this code to the tool. \\\n",
    "Another example: if a user ask 'What is log(5)' generate 'import math; interpreter_result = math.log(5)'\"\n",
    "\n",
    "def chat(message, history):\n",
    "    messages = [{\"role\": \"system\", \"content\": system_message}] + history + [{\"role\": \"user\", \"content\": message}]\n",
    "    response = openai.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "\n",
    "    if response.choices[0].finish_reason==\"tool_calls\":\n",
    "        message = response.choices[0].message\n",
    "        print(message)\n",
    "        response = handle_tool_call(message)\n",
    "        print(response)\n",
    "        messages.append(message)\n",
    "        messages.append(response)\n",
    "        response = openai.chat.completions.create(model=MODEL, messages=messages)\n",
    "    \n",
    "    return response.choices[0].message.content\n",
    "\n",
    "\n",
    "def handle_tool_call(message):\n",
    "    tool_call = message.tool_calls[0]\n",
    "    arguments = json.loads(tool_call.function.arguments)\n",
    "    my_code = arguments.get('my_code')\n",
    "    interpreter_result = evaluate_math_expression(my_code)\n",
    "    response = {\n",
    "        \"role\": \"tool\",\n",
    "        \"content\": json.dumps({\"my_code\": my_code,\"interpreter_result\": interpreter_result}),\n",
    "        \"tool_call_id\": tool_call.id\n",
    "    }\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3e50093-d7b6-4972-a8ba-6964f22218d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "gr.ChatInterface(fn=chat, type=\"messages\").launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75c81d73-d2d6-4e6b-8511-94d4a725f595",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
